﻿namespace DesktopRibbon.IconCache
{
    using System;
    using System.ComponentModel;
    using System.Data;
    using System.Data.SQLite;
    using System.IO;
    using System.Net;
    using System.Windows;
    using System.Windows.Media;
    using System.Windows.Media.Imaging;

    /// <summary>
    /// The icon cache loads icons from various sources in the background and provides cached values.
    /// </summary>
    class IconCache : IDisposable
    {
        #region Fields

        SQLiteConnection sqliteConnection;

        #endregion

        #region Constructors

        public IconCache()
        {
            string dbFile = App.DataDir() + "\\iconcache-v2.sqlite";
            bool dbInit = ! File.Exists(dbFile);
            sqliteConnection = new SQLiteConnection(
                string.Format(
                    "Data Source={0};",
                    dbFile.ToString()));
            try
            {
                sqliteConnection.Open();
                if (dbInit)
                {
                    using (var cmd = new SQLiteCommand("create table iconcache (path text, smallicon blob, largeicon blob, primary key(path));", sqliteConnection))
                    {
                        cmd.ExecuteNonQuery();
                    }
                }
            }
            catch (Exception)
            {
                sqliteConnection.Close();
            }
        }

        #endregion

        #region Delegates

        /// <summary>
        /// The IconLoadedHandler is used in the IconLoaded event.
        /// </summary>
        /// <param name="sender">The IconCache object that loaded the icon.</param>
        /// <param name="args">The loaded icon and the path that the icon has been loaded from.</param>
        public delegate void IconLoadedHandler(IconCache sender, IconLoadedEventArgs args);

        #endregion

        #region Events

        /// <summary>
        /// IconLoaded is emitted when LoadIcon has finished loading the icon in the background.
        /// </summary>
        public event IconLoadedHandler IconLoaded;

        #endregion


        #region Methods

        /// <summary>
        /// Closes the sqlite connection.
        /// </summary>
        public void Dispose()
        {
            sqliteConnection.Dispose();
        }

        /// <summary>
        /// LoadIcon only returns an ImageSource when an image exists in the cache. Otherwise,
        /// it will return null and start the icon loading in a background thread. When the icon
        /// has successfully been loaded, the IconLoaded event is emitted. The IconLoaded handler will be
        /// called in the main thread, not in the worker thread.
        /// </summary>
        /// <param name="path">The path that the icon should be loaded from.</param>
        /// <param name="userData">userData will be forwarded to the IconLoaded event and can be used to identify the loaded icon.</param>
        /// <returns>Either an cached icon, or null.</returns>
        public ImageSource LoadIcon(string path, object userData = null)
        {
            if (sqliteConnection.State == ConnectionState.Open)
            {
                using (var cmd = new SQLiteCommand("select smallicon, largeicon from iconcache where path=@path;", sqliteConnection))
                {
                    cmd.Parameters.Add(new SQLiteParameter("@path", DbType.String, 1024*4) { Value = path });
                    using (var reader = cmd.ExecuteReader())
                    {
                        if (reader.HasRows)
                        {
                            do
                            {
                                ImageSource smallIcon = LoadIconFromDbValue(reader.GetValue(0));
                                ImageSource largeIcon = LoadIconFromDbValue(reader.GetValue(1));
                                IconLoaded(this, new IconLoadedEventArgs(path, userData, smallIcon, largeIcon, false));
                            } while (reader.NextResult());
                        }
                    }
                }
            }

            var worker = new BackgroundWorker();
            worker.DoWork += LoadIconWorker;
            worker.RunWorkerCompleted += LoadIconWorkerCompleted;
            worker.RunWorkerAsync(Tuple.Create(path, userData));
            return null;
        }

        private ImageSource LoadIconFromDbValue(object dbValue)
        {
            object iconField = dbValue;
            if (! iconField.GetType().Equals(typeof(System.DBNull)))
            {
                var stream = new MemoryStream();
                var buffer = (byte[])iconField;
                stream.Write(buffer, 0, buffer.Length);
                stream.Seek(0, SeekOrigin.Begin);
                return LoadIconFromStream(stream);
            }
            return null;
        }

        private class LoadIconWorkerArgs
        {
            public string path;
            public Stream smallIconStream;
            public Stream largeIconStream;
            public object userData;
            public bool isImageFile;
        }

        private void LoadIconWorkerCompleted(object sender, RunWorkerCompletedEventArgs e)
        {
            System.Diagnostics.Trace.Write(string.Format("LoadIconWorkerCompleted id: {0}\n", System.Threading.Thread.CurrentThread.ManagedThreadId));
            if (e.Result == null)
            {
                return;
            }

            var args = (LoadIconWorkerArgs)e.Result;
            if (args.smallIconStream == null && args.largeIconStream == null)
            {
                // IconLoaded will not be called, when the load failed.
                return;
            }

            // First, read the whole stream in a byte[] and a MemoryStream so that the BitmapDecoder can decode that stream again.
            // We can not re-use the same stream, because some streams do not support Seek (e.g. streams that are a result of a WebRequest for favicons).

            var smallIconData = ReadIconFromStream(args.smallIconStream);
            if (args.smallIconStream == args.largeIconStream)
                args.largeIconStream.Seek(0, SeekOrigin.Begin);
            var largeIconData = ReadIconFromStream(args.largeIconStream);
            byte [] smallIconBytes = null;
            if (smallIconData != null)
                smallIconBytes = smallIconData.Item1;
            byte [] largeIconBytes = null;
            if (largeIconData != null)
                largeIconBytes = largeIconData.Item1;
            SaveIconsToCache(args.path, smallIconBytes, largeIconBytes);

            ImageSource smallIcon = null;
            ImageSource largeIcon = null;
            if (smallIconData != null)
                smallIcon = LoadIconFromStream(smallIconData.Item2);
            if (largeIconData != null)
                largeIcon = LoadIconFromStream(largeIconData.Item2);
            var handler = IconLoaded;
            if (null != handler)
            {
                handler(this, new IconLoadedEventArgs(args.path, args.userData, smallIcon, largeIcon, args.isImageFile));
            }
        }

        private static Tuple <byte [], Stream> ReadIconFromStream (Stream stream)
        {
            if (stream == null)
                return null;
            var iconStream = new MemoryStream();
            const int bufferSize = 1024;
            var membuffer = new byte[bufferSize];
            int readbytes;
            while ((readbytes = stream.Read(membuffer, 0, bufferSize)) > 0)
            {
                iconStream.Write(membuffer, 0, readbytes);
            }
            iconStream.Seek(0, SeekOrigin.Begin);
            return Tuple.Create(iconStream.ToArray(), (Stream) iconStream);
        }

        private ImageSource LoadIconFromStream(Stream memstream)
        {
            if (memstream == null || memstream.Length == 0)
                return null;

            var decoder = BitmapDecoder.Create(memstream, BitmapCreateOptions.None, BitmapCacheOption.OnLoad);
            if (decoder.Frames.Count > 0)
            {
                var icon = decoder.Frames[0];
                if (icon.Width > 1 && icon.Height > 1)
                {
                    return icon;
                    
                }
            }
            return null;
        }

        /// <summary>
        /// Saves the icon to the SQLite cache storage.
        /// </summary>
        /// <param name="path">The path that the icon has been loaded for.</param>
        /// <param name="smallIconData">The raw small icon data</param>
        /// <param name="largeIconData">The raw large icon data</param>
        private void SaveIconsToCache(string path, byte [] smallIconData, byte [] largeIconData)
        {
            try
            {
                using (var cmd = new SQLiteCommand("delete from iconcache where path=@path;", sqliteConnection))
                {
                    cmd.Parameters.Add(new SQLiteParameter("@path", DbType.String, 1024*4) { Value = path });
                    cmd.ExecuteNonQuery();
                }
                using (var cmd = new SQLiteCommand("insert into iconcache(path, smallicon, largeicon) values (@path, @smallicon, @largeicon);", sqliteConnection))
                {
                    cmd.Parameters.Add(new SQLiteParameter("@path", DbType.String, 1024*4) { Value = path });
                    cmd.Parameters.Add(new SQLiteParameter("@smallicon", DbType.Binary) { Value = smallIconData });
                    cmd.Parameters.Add(new SQLiteParameter("@largeicon", DbType.Binary) { Value = largeIconData });
                    cmd.ExecuteNonQuery();
                }
            }
            catch(Exception)
            {
            }
        }

        private void LoadIconWorker(object sender, DoWorkEventArgs e)
        {
            System.Diagnostics.Trace.Write(string.Format("LoadIconWorker id: {0}\n", System.Threading.Thread.CurrentThread.ManagedThreadId));

            var args = (Tuple <string, object>) e.Argument;
            string path = args.Item1;
            object userData = args.Item2;
            Stream smallIconStream = null;
            Stream largeIconStream = null;
            bool isImageFile = false;

            string pathUpper = path.ToUpper();
            if (pathUpper.StartsWith("HTTP:") || pathUpper.StartsWith("HTTPS:"))
            {
                var webRequest = WebRequest.Create("http://google.com/favicon.ico");
                var response = (HttpWebResponse) webRequest.GetResponse();
                smallIconStream = response.GetResponseStream();
                var memstream = new MemoryStream();
                byte [] bytes = new byte[1024];
                while (smallIconStream.Read(bytes, 0, 1024) > 0) {
                    memstream.Write(bytes, 0, 1024);
                }
                memstream.Seek(0, System.IO.SeekOrigin.Begin);
                smallIconStream = memstream;
            }
            else
            {
                if (pathUpper.EndsWith(".DLL") || pathUpper.EndsWith(".EXE"))
                {
                    var iconExtract = new TKageyu.Utils.IconExtractor(path);
                    if (iconExtract.IconCount > 0)
                    {
                        var exeIcon = iconExtract.GetIcon(0);
                        var icon = System.Windows.Interop.Imaging.CreateBitmapSourceFromHIcon(exeIcon.Handle, Int32Rect.Empty, BitmapSizeOptions.FromEmptyOptions());
                        var encoder = new PngBitmapEncoder();
                        encoder.Frames.Add(BitmapFrame.Create(icon));
                        smallIconStream = MemoryStreamFromEncoder(encoder);
                        largeIconStream = smallIconStream;
                    }
                }

                try {
                    var encoder = new PngBitmapEncoder();
                    var frame = BitmapFrame.Create(new Uri(path));
                    encoder.Frames.Add(frame);
                    largeIconStream = MemoryStreamFromEncoder(encoder);
                    isImageFile = true;
                } catch (Exception)
                {   
                }

                // Fallback: Load shell icon
                if (smallIconStream == null && largeIconStream == null)
                {
                    var smallIcon = ThirdParty.ShellIcon.GetSmallIcon(path);
                    var largeIcon = ThirdParty.ShellIcon.GetSmallIcon(path);
                    smallIconStream = MemoryStreamFromIcon(smallIcon);
                    largeIconStream = MemoryStreamFromIcon(largeIcon);
                }
            }
            e.Result = new LoadIconWorkerArgs() { 
                path = path,
                userData = userData,
                smallIconStream = smallIconStream,
                largeIconStream = largeIconStream,
                isImageFile = isImageFile
            };
        }

        private MemoryStream MemoryStreamFromEncoder(BitmapEncoder encoder)
        {
            var result = new MemoryStream();
            encoder.Save(result);
            result.Seek(0, SeekOrigin.Begin);
            return result;
        }

        private MemoryStream MemoryStreamFromIcon(System.Drawing.Icon icon)
        {
            var result = new MemoryStream();
            icon.Save(result);
            result.Seek(0, SeekOrigin.Begin);
            return result;
        }

        #endregion

    }
}
